#include "kernel_operator.h"

constexpr int32_t N = 1;
constexpr int32_t C = 1;
constexpr int32_t H = 8;
constexpr int32_t W = 256;

// 卷积核参数
constexpr int32_t kH = 2;
constexpr int32_t kW = 2;
constexpr int32_t stride = 1;
constexpr int32_t padding = 0;

// 计算输入尺寸（col矩阵）
constexpr int32_t out_H = (H + 2 * padding - kH) / stride + 1;  
constexpr int32_t out_W = (W + 2 * padding - kW) / stride + 1; 
constexpr int32_t L = out_H * out_W; 
constexpr int32_t input_channels = C * kH * kW;  // col矩阵的通道数

constexpr int32_t INPUT_LENGTH = N * input_channels * L;  // col矩阵大小
constexpr int32_t OUTPUT_LENGTH = N * C * H * W;  // 输出图像大小
constexpr int32_t USE_CORE_NUM = 8;
constexpr int32_t BLOCK_LENGTH = OUTPUT_LENGTH / USE_CORE_NUM;
constexpr int32_t TILE_NUM = 8;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM;

class KernelCol2Im {
public:
    __aicore__ inline KernelCol2Im() {}
    
    __aicore__ inline void Init(GM_ADDR col, GM_ADDR img, 
                               int32_t kernel_h, int32_t kernel_w, 
                               int32_t stride_val, int32_t padding_val)
    {
        this->kernel_h = kernel_h;
        this->kernel_w = kernel_w;
        this->stride_val = stride_val;
        this->padding_val = padding_val;
        
        // 计算col矩阵的尺寸
        this->out_H = (H + 2 * this->padding_val - this->kernel_h) / this->stride_val + 1;
        this->out_W = (W + 2 * this->padding_val - this->kernel_w) / this->stride_val + 1;
        this->L = this->out_H * this->out_W;
        this->input_channels = C * this->kernel_h * this->kernel_w;

        // 输入col矩阵和输出图像的GM地址
        colGm.SetGlobalBuffer((__gm__ float *)col, INPUT_LENGTH);
        imgGm.SetGlobalBuffer((__gm__ float *)img + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
        
        // 初始化缓冲区
        pipe.InitBuffer(outQueueImg, BUFFER_NUM, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(accumBuffer, TILE_LENGTH * sizeof(float));
        pipe.InitBuffer(tempBuffer, TILE_LENGTH * sizeof(float));
    }
    
    __aicore__ inline void Process()
    {
        int32_t loopCount = TILE_NUM * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void Compute(int32_t progress)
    {
        AscendC::LocalTensor<float> imgLocal = outQueueImg.AllocTensor<float>();
        AscendC::LocalTensor<float> accumLocal = accumBuffer.Get<float>();
        AscendC::LocalTensor<float> tempLocal = tempBuffer.Get<float>();

        // 计算当前tile在输出图像中的起始位置
        int32_t outputStart = AscendC::GetBlockIdx() * BLOCK_LENGTH + progress * TILE_LENGTH;
        
        // 初始化累加器为0
        AscendC::Duplicate(accumLocal, 0.0f, TILE_LENGTH);
        
        // 遍历所有kernel位置进行批量累加
        for (int32_t kh = 0; kh < kernel_h; kh++) {
            for (int32_t kw = 0; kw < kernel_w; kw++) {
                // 初始化临时buffer
                AscendC::Duplicate(tempLocal, 0.0f, TILE_LENGTH);
                
                // 批量收集当前kernel位置对所有像素的贡献
                for (int32_t i = 0; i < TILE_LENGTH; i++) {
                    int32_t outputIdx = outputStart + i;
                    
                    // 边界检查
                    if (outputIdx >= OUTPUT_LENGTH) {
                        continue;
                    }
                    
                    // 解析像素坐标
                    int32_t n_idx, c_idx, h_idx, w_idx;
                    PixelIndexToNCHW(outputIdx, n_idx, c_idx, h_idx, w_idx);
                    
                    // 检查该kernel位置是否对当前像素有贡献
                    int32_t h_out, w_out;
                    if (!CalculateOutputPosition(h_idx, w_idx, kh, kw, h_out, w_out)) {
                        continue;
                    }
                    
                    // 计算col索引
                    int32_t colIdx = CalculateColIndex(n_idx, c_idx, kh, kw, h_out, w_out);
                    
                    if (colIdx >= 0 && colIdx < INPUT_LENGTH) {
                        float value = colGm.GetValue(colIdx);
                        tempLocal.SetValue(i, value);
                    }
                }
                
                // 批量累加：accumLocal += tempLocal
                AscendC::Add(accumLocal, accumLocal, tempLocal, TILE_LENGTH);
            }
        }
        
        // 复制结果到输出队列
        AscendC::DataCopy(imgLocal, accumLocal, TILE_LENGTH);
        
        tempBuffer.FreeTensor(tempLocal);
        accumBuffer.FreeTensor(accumLocal);
        outQueueImg.EnQue<float>(imgLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
        AscendC::LocalTensor<float> imgLocal = outQueueImg.DeQue<float>();
        
        int32_t outputOffset = progress * TILE_LENGTH;
        int32_t copyLength = TILE_LENGTH;
        
        // 边界检查
        int32_t remaining = BLOCK_LENGTH - outputOffset;
        if (copyLength > remaining) {
            copyLength = remaining;
        }
        
        if (copyLength > 0) {
            AscendC::DataCopy(imgGm[outputOffset], imgLocal, copyLength);
        }
        
        outQueueImg.FreeTensor(imgLocal);
    }

    // 将像素索引转换为NCHW坐标
    __aicore__ inline void PixelIndexToNCHW(int32_t pixelIdx, int32_t& n, int32_t& c, int32_t& h, int32_t& w)
    {
        int32_t HW = H * W;
        int32_t CHW = C * HW;
        
        n = pixelIdx / CHW;
        int32_t remaining = pixelIdx % CHW;
        c = remaining / HW;
        remaining = remaining % HW;
        h = remaining / W;
        w = remaining % W;
    }

    // 计算输出位置
    __aicore__ inline bool CalculateOutputPosition(int32_t h, int32_t w, int32_t kh, int32_t kw, 
                                                     int32_t& h_out, int32_t& w_out)
    {
        int32_t h_numerator = h + padding_val - kh;
        int32_t w_numerator = w + padding_val - kw;
        
        // 检查是否能被stride整除
        if (h_numerator % stride_val != 0 || w_numerator % stride_val != 0) {
            return false;
        }
        
        h_out = h_numerator / stride_val;
        w_out = w_numerator / stride_val;
        
        // 检查输出位置是否在有效范围内
        if (h_out < 0 || h_out >= out_H || w_out < 0 || w_out >= out_W) {
            return false;
        }
        
        return true;
    }

    // 计算col矩阵索引
    __aicore__ inline int32_t CalculateColIndex(int32_t n, int32_t c, int32_t kh, int32_t kw, 
                                                  int32_t h_out, int32_t w_out)
    {
        int32_t channel_block = c * (kernel_h * kernel_w) + kh * kernel_w + kw;
        int32_t pos_idx = h_out * out_W + w_out;
        int32_t colIdx = n * (input_channels * L) + channel_block * L + pos_idx;
        
        return colIdx;
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueImg;
    AscendC::TBuf<AscendC::QuePosition::VECCALC> accumBuffer;
    AscendC::TBuf<AscendC::QuePosition::VECCALC> tempBuffer;
    AscendC::GlobalTensor<float> colGm;
    AscendC::GlobalTensor<float> imgGm;

    int32_t kernel_h;
    int32_t kernel_w;
    int32_t stride_val;
    int32_t padding_val;
    int32_t L;
    int32_t input_channels;
    int32_t out_H;
    int32_t out_W;
};

extern "C" __global__ __aicore__ void col_2_im_custom(GM_ADDR col, GM_ADDR img, 
                                                   int32_t kernel_h, int32_t kernel_w,
                                                   int32_t stride_val, int32_t padding_val)
{
    KernelCol2Im op;
    op.Init(col, img, kernel_h, kernel_w, stride_val, padding_val);
    op.Process();
}

#ifndef ASCENDC_CPU_DEBUG
void col_2_im_custom_do(uint32_t blockDim, void *stream, 
                      uint8_t *col, uint8_t *img,
                      int32_t kernel_h, int32_t kernel_w,
                      int32_t stride_val, int32_t padding_val)
{
    col_2_im_custom<<<blockDim, nullptr, stream>>>(col, img, kernel_h, kernel_w, stride_val, padding_val);
}
#endif